import os
import sys
import time
import yaml

import lightning as L
import pytorch_lightning as pl

from data_transform import *
from networks import *
from utils import get_args,class_builder
from pytorch_lightning.callbacks import ModelCheckpoint


def load_config(config_path):
    with open(config_path, 'r') as file:
        return yaml.safe_load(file)

def set_tb(model_type):
    Timestamp = time.strftime('_%m%d_%H_%M_%S')

    writer_path = os.path.join('logs', model_type, 'A'+ Timestamp)

    if not os.path.exists(writer_path):
        os.makedirs(writer_path)

    return Timestamp, writer_path


def main():
    args = get_args()
    config = load_config(args.config_path)

    Timestamp, writer_path= set_tb(model_type = config['network']['type'])
    
    # 日志
    logger = class_builder(config['Logger'], save_dir = writer_path)

    # TO DO:
    # 数据模块
    data_module = class_builder(config['Data'])
    data_module.setup()

    args.model_type = config['network']['type']
    print(args)

    print('train dataset length', len(data_module.train_dataloader().dataset))
    print('val dataset length', len(data_module.val_dataloader().dataset))

    # 训练器
    early_stop_callback = L.pytorch.callbacks.EarlyStopping(
        monitor="val_loss", 
        mode="min", 
        patience=50)

    ckpt_callback = L.pytorch.callbacks.ModelCheckpoint(
        monitor='val_loss',
        save_top_k=1,
        mode='min',
    ) 

    # 设置随机数种子
    L.fabric.seed_everything(config['Seed'])
    trainer = L.Trainer(**config['Trainer']['args'], 
                        logger=logger,
                        callbacks=[early_stop_callback, ckpt_callback])

    model = class_builder(
        config['network'],
        normalizer_y=data_module.normalizer_y, 
        normalizer_x=data_module.normalizer_x,
        use_wandb = False,)

        
    summary = pl.utilities.model_summary.ModelSummary(model,max_depth=1)
    print(summary) 
        
    # 训练
    trainer.fit(model=model, 
                train_dataloaders=data_module.train_dataloader(), 
                val_dataloaders=data_module.val_dataloader())
    
        
if __name__ == '__main__':
    main()
